import math

from myseq import MySeq
from MyMotifs import MyMotifs, print_matrix
from MotifFinding import MotifFinding

class MyMotifs_ex(MyMotifs):

    def __init__(self,seqs = [], pwm = [], alphabet = None):
        super().__init__(seqs,pwm,alphabet)
        self.__bulid_alphabet_dict()

    def __bulid_alphabet_dict(self):
        self.alphabet_dict = dict()
        for v,k in enumerate(self.alphabet):
            self.alphabet_dict[k] = v

    def add_pseudo_count_pwm(self,pseudo):
        for i in range(len(self.alphabet)):
            for j in range(self.size):
                self.pwm[i][j] += float(pseudo)/len(self.seqs)
    
    def score_sequence(self, seq):
        score = 0
        for i in range(self.size):
            lin = self.alphabet_dict[seq[i]]
            score += math.log10(self.pwm[lin][i])
        return score
    
    def high_score_motif(self,seq,threshold):
        res = list()
        for k in range(len(seq)-self.size):
            sub_seq = seq[k:k+ self.size]
            if self.score_sequence(sub_seq) >= threshold:
                res.append(sub_seq)
        return res
    
    def create_pwm_bg(self,bg_freq=0):
        self.create_pwm()
        for j in range(self.size):
            count_bg = 0
            W = 0
            update_bg_index = list()
            for i in range(len(self.alphabet)):
                if self.counts[i][j] == 0:
                    self.pwm[i][j] = bg_freq
                    count_bg += bg_freq
                else:
                    W += self.counts[i][j]
                    update_bg_index.append(i)
            for i in update_bg_index:
                self.pwm[i][j] = float(self.counts[i][j]) / len(self.seqs)
                self.pwm[i][j] = self.pwm[i][j] - self.counts[i][j]/W *count_bg
        print_matrix(self.pwm)

class MotifFinding_ex(MotifFinding):
        
    def create_motif_from_indexes(self, indexes):
        pseqs = []
        for i,ind in enumerate(indexes):
            pseqs.append(MySeq(seq=self.seqs[i][ind:(ind+self.motif_size)],
                                seq_type=self.seqs[i].get_seq_biotype()) )
        return MyMotifs_ex(pseqs)
    
    def heuristic_stochastic(self,fraction=0):
        if fraction < 0 or fraction > 1:
            raise ValueError('background frequerice fraction in (0,1]')
        
        from random import randint
        bg_freq = fraction/(len(self.alphabet)*self.motif_size)
        
        s = [0]* len(self.seqs) 
        for k in range(len(s)):
            s[k] = randint(0, self.seq_size(k)- self.motif_size)
        
        motif = self.create_motif_from_indexes(s)
        motif.create_pwm_bg(bg_freq=bg_freq)
        sc = self.score_multiplicative(s)
        bestsol = s
        improve = True
        while(improve):
            for k in range(len(s)):
                s[k] = motif.most_probable_sequence(self.seqs[k])
            if self.score_multiplicative(s) > sc: 
                sc = self.score_multiplicative(s)
                bestsol = s
                motif = self.create_motif_from_indexes(s)
                motif.create_pwm()
            else:
                improve = False    
        return bestsol


def test_psuedo_count():
    seq1 = MySeq("AAAGTT")
    seq2 = MySeq("CACGTG")
    seq3 = MySeq("TTGGGT")
    seq4 = MySeq("GACCGT")
    seq5 = MySeq("AACCAT")
    seq6 = MySeq("AACCCT")
    seq7 = MySeq("AAACCT")
    seq8 = MySeq("GAACCT")
    lseqs = [seq1, seq2, seq3, seq4, seq5, seq6, seq7, seq8]
    motifs = MyMotifs_ex(lseqs)
    print("raw pwm:")
    print_matrix(motifs.pwm)
    motifs.add_pseudo_count_pwm(0.5)
    print("add pseudo-count:")
    print_matrix(motifs.pwm)
    return motifs
    # print(motifs.alphabet_dict)

def test_high_score_motif(motif:MyMotifs_ex):
    print("consensus:",motif.consensus())
    print(motif.high_score_motif("CTATAAACCTTACATC",-3))

def test_background_frequencies():
    mf = MotifFinding_ex()
    mf.read_file("exampleMotifs.txt","DNA")
    print("Heuristic stochastic")
    sol = mf.heuristic_stochastic(1)
    print ("Solution: " , sol) #[1, 4, 45, 5, 0] 34
    print ("Score:" , mf.score(sol))
    print ("Score mult:" , mf.score_multiplicative(sol))
    print("Consensus:", mf.create_motif_from_indexes(sol).consensus())

if __name__ == "__main__":
    import os
    path = "/bioinformatic_algorithm_colg/8.Probabilistic_Motifs_and_Stochastic_Algorithms"
    os.chdir(os.getcwd()+path)
    print("### ex1 ###")
    motif = test_psuedo_count()
    print("### ex2 ###")
    test_high_score_motif(motif)
    print("### ex3 ###")
    test_background_frequencies()